package com.kingsoft.dc.khaos.plugin.core.utils;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLUseStatement;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.stat.TableStat.Name;
import com.alibaba.druid.util.JdbcConstants;
import com.ksyun.kbdp.dts.job.core.kit.Table;
import com.ksyun.kbdp.dts.job.core.kit.TableEdge;
import org.apache.commons.lang3.StringUtils;

import java.util.*;

/**
 * sql解析工具类.
 */
public class SqlParser {

    /**
     * 提取sql中table级别血缘关系
     * @param sql 待转换的sql，支持多条sql解析
     * @param dsId 数据源ID
     * @return
     */
    public static List<TableEdge> getTableEdgeList(String sql, String dsId) {
        List<TableEdge> tableEdgeList = new ArrayList<>();
        List<Map<String, TreeSet<String>>> fromToList = getFromTo(sql);
        fromToList.forEach(fromTo -> {
            fromTo.forEach((to, fromSet) -> {
                Table toTable = buildTable(to, dsId);
                if (fromSet.isEmpty()) {
                    tableEdgeList.add(TableEdge.newTableEdge(new Table(), toTable));
                } else {
                    fromSet.forEach(from -> {
                        Table fromTable = buildTable(from, dsId);
                        tableEdgeList.add(TableEdge.newTableEdge(fromTable, toTable));
                    });
                }

            });
        });
        return tableEdgeList;
    }

    /**
     * 构建Table实例
     * @param fullName 表全名,如db.table
     * @param dsId 数据源ID
     * @return
     */
    public static Table buildTable(String fullName, String dsId) {
        String[] splits = StringUtils.split(fullName, ".");
        Table table = new Table();
        table.setDsId(dsId);
        table.setDbName(splits[0]);
        table.setTableName(splits[1]);
        return table;
    }


    /**
     * 提取sql中血缘关系
     * @param sql 待分析的sql，支持多条sql
     * @return
     * @throws ParserException
     */
    public static List<Map<String, TreeSet<String>>> getFromTo(String sql) throws ParserException {
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, JdbcConstants.HIVE);
        if (stmts == null) {
            return null;
        }
        List<Map<String, TreeSet<String>>> fromToList = new ArrayList<>();
        String database = "DEFAULT";
        for (SQLStatement stmt : stmts) {
            SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(JdbcConstants.HIVE);
            if (stmt instanceof SQLUseStatement) {
                database = ((SQLUseStatement) stmt).getDatabase().getSimpleName();
                continue;
            }
            TreeSet<String> fromSet = new TreeSet<>();
            TreeSet<String> toSet = new TreeSet<>();
            stmt.accept(statVisitor);
            Map<Name, TableStat> tables = statVisitor.getTables();
            if (tables != null) {
                final String db = database;
                tables.forEach((tableName, stat) -> {
                    if (stat.getCreateCount() > 0 || stat.getInsertCount() > 0 || stat.getDeleteCount() > 0 || stat.getUpdateCount() > 0) {
                        String to = tableName.getName();
                        if (!to.contains("."))
                            to = db + "." + to;
                        toSet.add(to);
                    } else if (stat.getSelectCount() > 0) {
                        String from = tableName.getName();
                        if (!from.contains("."))
                            from = db + "." + from;
                        fromSet.add(from);
                    }
                });
            }
            Map<String, TreeSet<String>> fromTo = new HashMap();
            for (String toName : toSet) {
                fromTo.put(toName, fromSet);
            }
            fromToList.add(fromTo);

        }
        return fromToList;
    }
    public static void main(String[] args) {
//        String sql = "insert into tbl1 (id ,name) select id,name from sel_tab1; into insert_tbl2(id,name) values (1,';');insert into insert_tbl3(id,name) values (1,',')";
        String sql0  = "use aa;";
        String sql1 = "insert into table db1.tablename1 select a,b,c FROM (select a, b, c from db2.tablename2) ali;";
        String sql2 = "insert into table db11.tablename11 VALUES (1, ';');";
        String sql3 = "update table1 set id = 1 where id in (select id from table2) and name in (select id from table3);";
        String sql4 = "update table1 set id = 1";
        List<String> sb = new ArrayList<>();
        sb.add(sql0);
        sb.add(sql1);
        sb.add(sql2);
        sb.add(sql3);
        sb.add(sql4);
        String sql = StringUtils.join(sb.toArray(), ";");
        List<TableEdge> tableEdgeList = getTableEdgeList(sql, "88");
        tableEdgeList.forEach(tableEdge -> {
            System.out.println("----------------");
            System.out.println(String.format("==>From:%s.%s.%s", tableEdge.getFromTable().getDsId(), tableEdge.getFromTable().getDbName(), tableEdge.getFromTable().getTableName()));
            System.out.println(String.format("==>To  :%s.%s.%s", tableEdge.getToTable().getDsId(), tableEdge.getToTable().getDbName(), tableEdge.getToTable().getTableName()));

        });

    }

}
